package encoder

import (
	"time"

	"github.com/pion/rtp"
	"github.com/yapingcat/gomedia/go-mpeg2"
	"gogs.smarteye.com/servere/gobvrtp/track"
)

// PSDecoder 解rtp包->解ps包
type PSDecoder struct {
	RTPDecoder

	demuxer *mpeg2.PSDemuxer
}

// Init initializes the decoder.
func (d *PSDecoder) Init() error {
	d.demuxer = mpeg2.NewPSDemuxer()

	return d.RTPDecoder.Init()
}

// Decode decodes an audio frame from a RTP packet.
func (d *PSDecoder) Decode(pkt *rtp.Packet, onFrame func(frame []byte, pts time.Duration, codec track.CodecId)) error {

	nalus, err := d.RTPDecoder.Decode(pkt)
	if err != nil {
		return err
	}

	d.demuxer.OnFrame = func(frame []byte, cid mpeg2.PS_STREAM_TYPE, pts uint64, dts uint64) {
		onFrame(frame, time.Duration(pts)*time.Millisecond, d.getCodecId(cid))
	}

	for _, nalu := range nalus {
		if err = d.demuxer.Input(nalu); err != nil {
			return err
		}
	}

	return nil
}

func (d *PSDecoder) getCodecId(t mpeg2.PS_STREAM_TYPE) track.CodecId {
	switch t {
	case mpeg2.PS_STREAM_AAC:
		return track.CodecIdAAC
	case mpeg2.PS_STREAM_H264:
		return track.CodecIdH264
	case mpeg2.PS_STREAM_H265:
		return track.CodecIdH265
	case mpeg2.PS_STREAM_G711A:
		return track.CodecIdG711A
	case mpeg2.PS_STREAM_G711U:
		return track.CodecIdG711U
	default:
		return track.CodecIdUnKnown
	}
}
